import math
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
from .cnn import BaseConv2d
from .norm import MaxMinNorm


class GradHead(nn.Module):
    def __init__(self, in_channels, kernel_size=1, activation=True):
        super().__init__()
        self.conv = BaseConv2d(in_channels, 2, kernel_size, act=activation)
        self.geps = math.sqrt(2) * 2e-2 # 5-pixel
        self.aoff = math.pi / 180
        self.norm = MaxMinNorm()
    
    def forward(self, x):
        b, c, h, w = x.size()

        x = self.norm(x)
        x = self.conv(x)
        gy, gx = x.split(1, dim=1)
        grad = torch.sqrt(gx*gx + gy*gy)
        grad = torch.clip(grad, 0, 1)
        angle = torch.atan2(gy, gx) # [-pi, pi]
        angle = (angle + math.pi  + self.aoff) / (2*math.pi) # [0+aeps, 2pi+apes]
        # print(angle.dtype, grad.dtype, type(self.geps))
        angle = torch.where(grad < self.geps, torch.zeros_like(angle), angle)
        return grad, angle
    
class MDD(nn.Module):
    # Main Direction Detect
    def __init__(self, in_channels, grid_size, stride=None, kernel_size=3, act=True, bins=9, topk=None):
        super().__init__()

        self.grid_size = grid_size
        self.stride = stride or grid_size
        self.bins = bins
        self.topk = topk or bins
        self.aoff = math.pi / 180
        self.grad = GradHead(in_channels, kernel_size, activation=act)
    
    def forward(self, x):
        grad, angle = self.grad(x)
        angle = torch.squeeze(angle, dim=1)
        b, h, w = angle.size()
        kh, kw = self.grid_size
        sh, sw = self.stride
        agl = angle.unfold(1, kh, sh).unfold(2, kw, sw).reshape(b, -1, kh*kw)
        l = agl.size(1) # windows size
        hist = torch.zeros(b, l, self.bins)
        for i in range(b):
            for j in range(l):
                jhist = torch.histc(agl[i, j], self.bins, min=self.aoff, max=1+self.aoff)
                hist[i, j] = jhist.float()

        hist, _ = torch.topk(hist, k=self.topk, dim=2)
        
        return hist, grad # b, n, c
        
    
class HOG(nn.Module):
    def __init__(self, in_channels, grid_size:tuple, stride:tuple=None, kernel_size=3, act=True, bins=9):
        super().__init__()

        self.grid_size = grid_size
        self.stride = stride or grid_size
        self.bins = bins
        self.grad = GradHead(in_channels, kernel_size, activation=act)
    
    def forward(self, x):
        grad, angle = self.grad(x)
        b, _, h, w = grad.size()
        kh, kw = self.grid_size
        sh, sw = self.stride
        step = math.pi / self.bins
        grad_maps = []
        for i in range(self.bins):
            gfeat = torch.where((angle > step*i) & (angle < step * (i+1)), grad, torch.zeros_like(grad))
            grad_maps.append(gfeat)
        
        grad_maps = torch.cat(grad_maps, dim=1) # [b, bins, h, w]
        grads = grad_maps.unfold(2, kh, sh).unfold(3, kw, sw).reshape(b, self.bins, -1, kh, kh)
        return torch.sum(grads, dim=2), grad

            

if __name__ == "__main__":
    x = torch.rand(2, 3, 8, 8)
    glayer = GradHead(3)
    start = time.time()
    grad, angle = glayer(x)
    #print(grad, grad.shape)
    # angle = angle.squeeze(1)
    print(angle.shape)
    print(angle)
    
    # y = glayer.extract_features(angle, 9, (14, 14))
    print('estimate: ', time.time() - start)
    # print(y.shape)


    # net = HOG(3, (4, 4), stride=(2, 2))
    # y = net(x)

    # print(y.shape)
    # print(y)
